package com.example.demo.controller;

import com.alibaba.fastjson.JSONObject;
import com.example.demo.common.Result;
import com.example.demo.verify.NoEmpty;
import com.github.xiaoymin.knife4j.annotations.ApiOperationSupport;
import com.github.xiaoymin.knife4j.annotations.ApiSupport;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiImplicitParam;
import io.swagger.annotations.ApiImplicitParams;
import io.swagger.annotations.ApiOperation;
import lombok.extern.slf4j.Slf4j;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

/**
 * SSE 默认支持断线重连，当网络断线或服务器停止时，前端会一直尝试重新连接
 * onCompletion，当后端连接断开时执行，即：连接已结束/完结/终结
 * 若后端使用多线程循环发送数据，要处理当出错时结束线程
 */

@RestController
@RequestMapping("/api/sse")
@Api(tags = "5、SSE通信管理")
@ApiSupport(author = "罗海林", order = 4)
@Validated
@Slf4j
public class SseController {

    private AtomicInteger counter = new AtomicInteger(0);
    private static Map<String, SseEmitter> sseList = new ConcurrentHashMap<>();

    @GetMapping(value = "/stream")
    public SseEmitter streamMessages(@RequestParam String id) {
        log.info("start建立连接");

        // 超时时间0表示不过期
        //SseEmitter emitter = new SseEmitter();
        SseEmitter emitter = new SseEmitter(0L);
        //SseEmitter emitter = new SseEmitter(3 * 1000L);

        /*new Thread(() -> {
            try {
                while (true) {
                    log.info("发送数据");

                    JSONObject js = new JSONObject();
                    js.put("id", id);
                    js.put("name", "罗海林");
                    js.put("info", "消息：" + counter.incrementAndGet());
                    emitter.send(js.toJSONString());

                    Thread.sleep(2000);
                }
            }
            catch (Exception e) {
                emitter.completeWithError(e);
                log.info("发送数据出错");
            }
            finally {
                log.info("发送数据线程结束~~~~~~");
            }
        }).start();*/

        // 设置回调函数
        emitter.onCompletion(completionCallBack(id));
        emitter.onTimeout(timeoutCallBack(id));
        emitter.onError(errorCallBack(id));

        // 往SSE连接列表中增加新的用户ID与emitter（若先前有残留的旧sse，先删除，再添加新的）
        insertSseList(id, emitter);

        log.info("end返回连接");

        return emitter;
    }

    @PostMapping("/sendMessage")
    @ApiImplicitParams({
            @ApiImplicitParam(name = "userId", value = "用户ID"),
            @ApiImplicitParam(name = "message", value = "消息内容")
    })
    @ApiOperation(value = "发送消息给指定用户", response = Result.class)
    @ApiOperationSupport(author = "罗海林", order = 1)
    public Result sendMessage(@NoEmpty(message = "用户ID不能为空") String userId,
                              @NoEmpty(message = "消息内容不能为空") String message) {
        Result res = new Result();

        if (sseList.containsKey(userId)) {
            try {
                JSONObject js = new JSONObject();
                js.put("id", userId);
                js.put("name", "罗海林");
                js.put("info", "消息：" + message);
                sseList.get(userId).send(js.toJSONString());
                res.ok();
            }
            catch (Exception e) {
                removeSseInfo(userId);
                res.setBusinessError("发送数据出错");
            }
        }
        else {
            res.setBusinessError("该用户已下线");
        }

        return res;
    }

    @PostMapping("/checkSseList")
    @ApiOperation(value = "检查SSE连接列表", response = Result.class)
    @ApiOperationSupport(author = "罗海林", order = 2)
    public Result checkSseList() {
        Result res = new Result();
        res.ok();
        res.put("list", sseList.toString());
        res.put("count", sseList.size());
        return res;
    }

    private void insertSseList(String userId, SseEmitter emitter) {
        removeSseInfo(userId);
        sseList.put(userId, emitter);
    }

    private void removeSseInfo(String userId) {
        if (sseList.containsKey(userId)) {
            SseEmitter emitter = sseList.remove(userId);
            if (emitter != null) {
                try {
                    emitter.complete();
                }
                catch (Exception ignored) { }
            }
        }
    }

    private Runnable completionCallBack(String userId) {
        return () -> {
            log.info("用户-{} 连接已终结", userId);

            // 将SSE连接列表中的相应用户ID与emitter删除
            removeSseInfo(userId);
        };
    }

    private Runnable timeoutCallBack(String userId) {
        return () -> {
            log.info("用户-{} 连接超时", userId);
            removeSseInfo(userId);
        };
    }

    private Consumer<Throwable> errorCallBack(String userId) {
        return (Throwable t) -> {
            log.info("用户-{} 连接异常：{}", userId, t.getMessage());
            removeSseInfo(userId);
        };
    }
}
